package main

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"strconv"
)

func trimPDF(inputPath string, outputPath string, startPage, endPage int) error {
	// 检查 pdftk 是否可用
	_, err := exec.LookPath("pdftk")
	if err != nil {
		return fmt.Errorf("未找到 pdftk 命令，请先安装 pdftk")
	}

	// 获取 PDF 总页数
	cmd := exec.Command("pdftk", inputPath, "dump_data")
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("无法获取 PDF 信息: %v", err)
	}

	// 解析总页数的简单实现（这里是一个基本实现，可能需要更健壮的解析）
	totalPages := 0
	for _, line := range string(output) {
		// 这里是一个简化的实现，实际情况可能更复杂
		if line == 'N' {
			totalPages++
		}
	}

	// 验证页面范围
	if startPage < 1 || endPage > totalPages || startPage > endPage {
		return fmt.Errorf("页面范围无效: 总页数 %d, 起始页 %d, 结束页 %d", totalPages, startPage, endPage)
	}

	// 使用 pdftk 提取页面
	cmd = exec.Command("pdftk", inputPath, "cat", 
		fmt.Sprintf("%d-%d", startPage, endPage), 
		"output", outputPath)
	
	err = cmd.Run()
	if err != nil {
		return fmt.Errorf("PDF 裁剪失败: %v", err)
	}

	return nil
}

func mergePDFs(inputPaths []string, outputPath string) error {
	// 检查 pdftk 是否可用
	_, err := exec.LookPath("pdftk")
	if err != nil {
		return fmt.Errorf("未找到 pdftk 命令，请先安装 pdftk")
	}

	// 准备 pdftk 命令参数
	args := append(inputPaths, "cat", "output", outputPath)
	cmd := exec.Command("pdftk", args...)
	
	err = cmd.Run()
	if err != nil {
		return fmt.Errorf("PDF 合并失败: %v", err)
	}

	return nil
}

func main() {
	if len(os.Args) < 4 {
		fmt.Println("用法:")
		fmt.Println("  裁剪: pdf-tool trim <输入文件> <输出文件> <起始页> <结束页>")
		fmt.Println("  合并: pdf-tool merge <输出文件> <输入文件1> <输入文件2> [更多输入文件...]")
		os.Exit(1)
	}

	command := os.Args[1]

	switch command {
	case "trim":
		if len(os.Args) != 6 {
			fmt.Println("裁剪用法: pdf-tool trim <输入文件> <输出文件> <起始页> <结束页>")
			os.Exit(1)
		}
		inputPath := os.Args[2]
		outputPath := os.Args[3]
		startPage, err := strconv.Atoi(os.Args[4])
		if err != nil {
			log.Fatalf("起始页无效: %v", err)
		}
		endPage, err := strconv.Atoi(os.Args[5])
		if err != nil {
			log.Fatalf("结束页无效: %v", err)
		}

		err = trimPDF(inputPath, outputPath, startPage, endPage)
		if err != nil {
			log.Fatalf("PDF 裁剪失败: %v", err)
		}
		fmt.Printf("成功将 %s 的第 %d 页到第 %d 页裁剪到 %s\n", inputPath, startPage, endPage, outputPath)

	case "merge":
		if len(os.Args) < 5 {
			fmt.Println("合并用法: pdf-tool merge <输出文件> <输入文件1> <输入文件2> [更多输入文件...]")
			os.Exit(1)
		}
		outputPath := os.Args[2]
		inputPaths := os.Args[3:]

		err := mergePDFs(inputPaths, outputPath)
		if err != nil {
			log.Fatalf("PDF 合并失败: %v", err)
		}
		fmt.Printf("成功将 %d 个 PDF 文件合并到 %s\n", len(inputPaths), outputPath)

	default:
		fmt.Println("未知命令。支持的命令:")
		fmt.Println("  trim  - 裁剪 PDF")
		fmt.Println("  merge - 合并 PDF")
		os.Exit(1)
	}
}
